from fastmcp import FastMCP,Context
from fastmcp.tools import tool
from typing import Dict, Any, List
from typing import Annotated
from pydantic import Field
from datetime import datetime
import os
from wxdump.db import DBHandler
from fastmcp.server.auth import BearerAuthProvider
import json
from fastmcp.server.dependencies import get_context
import requests
from urllib.parse import quote


# auth = BearerAuthProvider(
#     jwks_uri="https://my-identity-provider.com/.well-known/jwks.json",
#     issuer="https://my-identity-provider.com/",
#     audience="my-mcp-server"
# )


mcp = FastMCP("WxDump")


def api_file(url):
    # 判断是否使用本地数据
    encoded_url = quote(url, safe='')
    
    # 根据条件构建基础URL
    base_url = "http:127.0.0.1:5000/api/rs/file"
    
    # 构建完整URL
    full_url = f"{base_url}?src={encoded_url}"
    
    # 发送GET请求
    response = requests.get(full_url)
    
    # 返回响应
    return response.content  # 或根据需要返回response.text、response.json()等
# 将 gc 的导入移到函数内部
def get_db():
    from wxdump.api.utils import gc
    my_wxid = gc.get_conf(gc.at, "last")
    db_config = gc.get_conf(my_wxid, "db_config")
    db = DBHandler(db_config, my_wxid=my_wxid)
    return db
def parse_messages(message,db,room_name):
    room_info = db.get_room_list(room_name)
    talker_name = "UnKonw"
    if room_info:
        wxid2userinfo = room_info.get(room_name).get("wxid2userinfo")
        talker_wxid = message.get("talker")

        if talker_wxid:
            user_info = wxid2userinfo.get(talker_wxid)
            if user_info:
                # 优先使用备注名，没有则使用昵称
                display_name = user_info.get("remark") or user_info.get("roomNickname")  or user_info.get("nickname")
                if display_name:
                    talker_name = display_name
    return talker_name
        # 1. 首先根据关键字查询联系人
def parse_time(time_str):
    """
    解析多种格式的时间字符串，支持以下格式：
    - YYYY-MM-DD
    - YYYY-MM-DD HH:MM 
    - YYYY-MM-DD HH:MM:SS
    
    返回datetime对象，可用于SQL查询
    """
    time_str = time_str.strip()
    
    # 尝试不同的时间格式
    formats = [
        "%Y-%m-%d",          # 2025-01-02
        "%Y-%m-%d %H:%M",    # 2025-01-02 18:20
        "%Y-%m-%d %H:%M:%S"  # 2025-01-02 18:02:20
        "%Y-%m-%d/%H:%M",    # 2025-01-02 18:20
        "%Y-%m-%d/%H:%M:%S",  # 2025-01-02 18:02:20
        "%Y-%m"
    ]
    
    # 处理可能出现的中文冒号
    time_str = time_str.replace('：', ':')
    
    for fmt in formats:
        try:
            return datetime.strptime(time_str, fmt).timestamp()
        except ValueError:
            continue
    
    # 如果所有格式都不匹配，抛出异常
    raise ValueError(f"不支持的时间格式: {time_str}，请使用 YYYY-MM-DD 或 YYYY-MM-DD HH:MM 或 YYYY-MM-DD HH:MM:SS 格式")

@mcp.tool
async def handle_query_contact( keyword: Annotated[str, Field(description="联系人的搜索关键词，可以是姓名、备注名或ID")]):

    '''

    功能描述:
    此工具用于查询用户的联系人信息，支持通过姓名、备注名或ID进行模糊查询。

    参数说明:
    - keyword (str): 联系人的搜索关键词，可以是姓名、备注名或ID。如果此参数为空，将返回所有联系人列表。

    使用场景:
    - 用户询问特定联系人的联系方式时。
    - 用户希望了解某个联系人的详细信息时。
    - 用户需要查找符合特定关键词的联系人时。

    返回值:
    (dict): 包含联系人信息的字典列表。每个字典包含 'user_name', 'alias', 'remark', 'nick_name' 等字段。

    '''
    try:
        db = get_db()
        user_list = db.get_user_list(word=keyword)
        result = "UserName,Account,Remark,NickName\n"
        for user in user_list.values():
            result += f"{user['wxid']},{user['account']},{user['remark']},{user['nickname']}\n"
        
        return {
            "content": [ result],
            "is_error": False
        }
    except Exception as e:
        return {
            "content": [{"type": "text", "text": f"查询联系人失败：{str(e)}"}],
            "is_error": True
        }


@mcp.tool
async def handle_query_recent_chat(
                            keyword: Annotated[str, Field(description="搜索关键词")],
                            start_time: Annotated[str|None, Field(description="开始时间，格式：YYYY-MM-DD HH:MM")] = None,
                            end_time: Annotated[str|None, Field(description="结束时间，格式：YYYY-MM-DD HH:MM")] = None):
    """

    功能描述:
    此工具用于查询用户的最近关于某个话题会话列表，包括个人聊天和群聊。

    参数说明:
    - keyword (str， 必填): 会话的搜索关键词，可以是会话名称或内容中的关键词。如果此参数为空，将返回所有最近会话列表。
    - start_time (str，选填): 查询的起始时间，格式为 "YYYY-MM-DD HH:MM"，如果不填则查询所有时间。
    - end_time (str，选填): 查询的结束时间，格式为 "YYYY-MM-DD HH:MM"，如果不填则查询所有时间。

    使用场景:
    - 用户希望通过关键词筛选最近会话时。

    返回值:
    (dict): 包含最近会话信息的字典列表。
    """
    try:
        parse_start_time = start_time
        parse_end_time = end_time
        if start_time:
            parse_start_time = parse_time(start_time)
        if end_time:
            parse_end_time = parse_time(end_time)
        db =  get_db()
        sessions = db.get_sessions(keyword, parse_start_time, parse_end_time)
        result = []
        for seesion in sessions:
            message_result = []
            for message in seesion["messages"]:
                talker_wxid  = message["talker"]
                if talker_wxid == "我":
                    # 查找对应的用户信息
                    talker_name = talker_wxid
                else:
                    talker_name = "Unknown"
                    room_name = message.get("room_name", "")
                    if room_name.endswith("@chatroom"):
                        talker_name = parse_messages(message,db, room_name)
        
                message_result.append({"talker": talker_name, "msg": message["msg"], "time": message["CreateTime"],
                            "msg": message["msg"],"src": message["src"]}) 
            user_list = db.get_user_list(word=seesion["talker"]) 
            user_info = user_list.get(seesion["talker"])
            result.append({"msgs":message_result,"talker":user_info["remark"] or user_info["nickname"] or user_info["account"], "time": seesion["time"]})

        
        return {
            "content": [{"type": "json", "text": result}],
            "is_error": False
        }
    except Exception as e:
        return {
            "content": [{"type": "text", "text": f"查询最近会话失败：{str(e)}"}],
            "is_error": True
        }

@mcp.tool
def handle_current_time():
    '''
    获取当前时间
    '''
    try:
        current_time = datetime.now().isoformat()
        return {
            "content": [{"type": "text", "text": current_time}],
            "is_error": False
        }
    except Exception as e:
        return {
            "content": [{"type": "text", "text": f"获取当前时间失败：{str(e)}"}],
            "is_error": True
        }


@mcp.tool
async def handle_query_messages(
    talker: Annotated[str, Field(description="可以是wxid,用户名、昵称、备注、描述，允许拼音")],
    user_type: Annotated[str|None, Field(description="user or chatroom")] =None,
    start_time: Annotated[str|None, Field(description="开始时间，格式：YYYY-MM-DD HH:MM")] = None,
    end_time: Annotated[str|None, Field(description="结束时间，格式：YYYY-MM-DD HH:MM")] = None,
    limit: Annotated[int, Field(description="返回结果的最大数量")] = 20):
    '''

    功能描述:
    此工具用于获取指定联系人的历史聊天记录。

    参数说明:
    - talker (str,必填): 查询关键字，可以是wxid,用户名、昵称、备注、描述，允许拼音。
    param user_type(str,选填): 联系人类型，None为所有，'contact'为个人联系人，'chatroom'为群聊
    - start_time (str，选填): 查询的起始时间，格式为 "YYYY-MM-DD HH:MM"，如果不填则查询所有时间。
    - end_time (str，选填): 查询的结束时间，格式为 "YYYY-MM-DD HH:MM"，如果不填则查询所有时间。
    - limit (int, optional): 返回结果的最大数量。默认为20，表示返回20条记录。

    使用场景:
    - 用户需要查找与某个联系人在特定时间段内的对话时
    - 用户想回顾与某个群组的历史交流内容时。

    返回值:
    (dict): 包含消息列表和用户信息列表的字典。消息列表包含每条消息的详细信息，如发送者、接收者、时间、内容等。
    示例:
    ```json
    {
        "content": [
            {
                "type": "text",
                "text": "消息1内容\n消息2内容\n..."
            }
        ],
        "is_error": false
    }
    ```
    '''
    try:
    # 时间格式转换
        parse_start_time = start_time
        parse_end_time = end_time
        if start_time:
            parse_start_time = parse_time(start_time)
        if end_time:
            parse_end_time = parse_time(end_time)
        db = get_db()
        user_list = db.get_user_list(word=talker,user_type=user_type)
        
        if not user_list:
            return {
            "content": [{"type": "text", "text": f"没有查到联系人"}],
            "is_error": True
        }
        
        # 2. 获取所有匹配联系人的wxid
        wxids = list(user_list.keys())
        # 3. 获取指定时间段内的聊天记录
        messages, talkers = db.get_msg_list(
            wxids=wxids,
            start_index=0,
            page_size=limit,
            start_createtime=parse_start_time,
            end_createtime=parse_end_time
        )
        # print(messages)
        result = []
        for message in messages:
            talker_wxid  = message["talker"]
            if talker_wxid == "我":
                # 查找对应的用户信息
                talker_name = talker_wxid
            else:
                talker_name = "Unknown"
                room_name = message.get("room_name", "")
                if room_name.endswith("@chatroom"):
                    talker_name = parse_messages(message,db, room_name)

                else:
                    for wxid, user in user_list.items():  # 修改这里，正确遍历字典
                        if wxid == talker_wxid:  # 比较wxid
                            talker_name = user.get('remark', user.get('account', ''))  # 使用get方法避免KeyError
                            break
                    else:
                        # 如果没有找到匹配的wxid，默认使用account
                        talker_name = user_list.get(talker_wxid, {}).get('account', 'Unknown')
    
            result.append({"talker": talker_name, "msg": message["msg"], "time": message["CreateTime"],
                        "msg": message["msg"],"src": message["src"],"type":message["type_name"]})  # 修改这里，使用更安全的访问方式# 修改这里，使用更安全的访问方式
        result_msg = "发送人,消息类型,消息内容,消息时间,附件地址\n"
        for msg in result:
            result_msg += f"{msg['talker']},{msg['type']},{msg['msg']},{msg['time']},{msg['src']}\n"
        return result_msg
    except Exception as e:
        return {
            "content": [{"type": "text", "text": f"查询消息失败：{str(e)}"}],
            "is_error": True
        }

@mcp.tool
async def handle_query_top_talkers(ctx:Context | None=None,
    top: Annotated[int, Field(description="返回前N个聊天最多的联系人")] = 10,
    start_time: Annotated[str|None, Field(description="开始时间，格式：YYYY-MM-DD")] =None,
    end_time: Annotated[str|None, Field(description="结束时间，格式：YYYY-MM-DD")] =None
):
    '''

    功能描述:
    此工具用于统计并返回聊天次数最多的联系人列表。

    参数说明:
    - top (int, optional): 返回聊天次数最多的前N个联系人。默认为10。
    - start_time (str, optional): 统计的开始时间，格式为 "YYYY-MM-DD"。如果为空，则不限制开始时间。
    - end_time (str, optional): 统计的结束时间，格式为 "YYYY-MM-DD"。如果为空，则不限制结束时间。

    使用场景:
    - 用户想知道自己最常聊天的联系人是谁时。
    - 用户需要分析某个时间段内（例如上周、上个月）的聊天活跃度时。
    - 用户想了解与特定联系人的聊天频率时。

    返回值:
    (dict): 包含聊天次数统计的联系人列表。每个联系人包含其名称和聊天次数。
    示例:
    ```json
    {
        "content": [
            {
                "type": "text",
                "text": "联系人名称,聊天次数\n张三,150\n李四,120\n王五,90\n"
            }
        ],
        "is_error": false
    }
    ```
    '''
    try:
        parse_start_time = start_time
        parse_end_time = end_time
        if start_time:
            parse_start_time = parse_time(start_time)
        if end_time:
            parse_end_time = parse_time(end_time)
        db =  get_db()
        # 创建数据库连接并获取统计数据
        talker_stats = db.get_top_talker_count(
            top=top,
            start_time=start_time,
            end_time=end_time
        )
        
        # 格式化返回结果
        if not talker_stats:
            return {
                "content": [{"type": "text", "text": "未找到聊天记录"}],
                "is_error": False
            }
            
        # 构建表格形式的输出
        result = "联系人,我发送次数,对方发送次数\n"
        for (key,value) in talker_stats.items():
            user_list = db.get_user_list(key)
            username = user_list[key]["remark"] or user_list[key]["nickname"] or user_list[key]["account"]
            result += f"{username},{value['sender_count']},{value['receiver_count']}\n"
        
        return {
            "content": [{"type": "text", "text": result}],
            "is_error": False
        }
        
    except Exception as e:
        return {
            "content": [{"type": "text", "text": f"获取聊天统计失败：{str(e)}"}],
            "is_error": True
        }
# Utility function that needs context but doesn't receive it as a parameter
async def process_data(data: list[float]) -> dict:
    # Get the active context - only works when called within a request
    ctx = get_context()    
    await ctx.info(f"Processing {len(data)} data points")
    

@mcp.prompt
async def get_promt():
    return """
    强制多步查询流程 (针对特定话题或发送者发言):
    当查询特定话题或特定发送者发言时，必须严格按照以下流程使用，任何偏离都会导致错误的结果：

    步骤1: 初步定位相关消息
    - 使用  handle_query_recent_chat 工具查找特定话题。
    - 使用 找到handle_query_recent_chat 返回话题的 ‘talker’ 和 ‘time’。
    - 使用 handle_query_messages 工具查找特定发送者，在'time'时间范围内的聊天消息。
    - 使用较宽时间范围初步查询。

    步骤2: 【必须执行】针对每个关键结果点分别获取上下文
    - 必须对步骤1返回的每个时间点T1, T2, T3...分别执行独立查询（时间范围接近的消息可以合并为一个查询）。
    - 每次独立查询必须移除 `keyword` 参数。
    - 每次独立查询使用 "Tn前后15-30分钟" 的窄范围。
    - 每次独立查询仅保留 `talker` 参数。

    步骤3: 【必须执行】综合分析所有上下文
    - 必须等待所有步骤2的查询结果返回后再进行分析。
    - 必须综合考虑所有上下文信息后再回答用户。

    严格执行规则:
    - 禁止仅凭步骤1的结果直接回答用户。
    - 禁止在步骤2使用过大的时间范围一次性查询所有上下文。
    - 禁止跳过步骤2或步骤3。
    - 必须对每个关键结果点分别执行独立的上下文查询。

    执行示例:
    正确流程示例:
    1. 步骤1: `chatlog(time="2023-04-01~2023-04-30", talker="工作群", keyword="项目进度")`
       返回结果: 4月5日、4月12日、4月20日有相关消息。
    2. 步骤2:
       - 查询1: `chatlog(time="2023-04-05/09:30~2023-04-05/10:30", talker="工作群")` // 注意没有keyword
       - 查询2: `chatlog(time="2023-04-12/14:00~2023-04-12/15:00", talker="工作群")` // 注意没有keyword
       - 查询3: `chatlog(time="2023-04-20/16:00~2023-04-20/17:00", talker="工作群")` // 注意没有keyword
    3. 步骤3: 综合分析所有上下文后回答用户。

    错误流程示例:
    - 仅执行步骤1后直接回答。
    - 步骤2使用 `time="2023-04-01~2023-04-30"` 一次性查询。
    - 步骤2仍然保留 `keyword` 或 `sender` 参数。

    自我检查 (回答用户前必须自问):
    - 我是否对每个关键时间点都执行了独立的上下文查询?
    - 我是否在上下文查询中移除了 `keyword` 和 `sender` 参数?
    - 我是否分析了所有上下文后再回答?
    - 如果上述任一问题答案为"否"，则必须纠正流程。

    返回值格式:
    - 单个 `talker` 查询: "昵称(ID) 时间\n消息内容\n昵称(ID) 时间\n消息内容"
    - 多个 `talker` 查询: "昵称(ID)\n[TalkerName(Talker)] 时间\n消息内容"

    重要提示:
    1. 当用户询问特定时间段内的聊天记录时，必须使用正确的时间格式，特别是包含小时和分钟的查询。
    2. 对于 "今天下午4点到5点聊了啥" 这类查询，正确的时间参数格式应为 start_time:"YYYY-MM-DD 16:00" end_time"YYYY-MM-DD 17:00"。

【单一时间点格式】
- 精确到日："2023-04-18"
- 精确到分钟（必须包含空格）："2023-04-18 14:30"（表示2023年4月18日14点30分）


【重要提示】包含小时分钟的格式必须使用空格和冒号：" "和":"
正确示例："2023-04-18 16:30"（4月18日下午4点30分）
错误示例："2023-04-18/16:30"、"2023-04-18T16:30"

- 可使用ID、昵称或备注名
    """
@mcp.tool
def user_session_list():
    """
    获取联系人列表
    :return:
    """
    db = get_db()
    ret = db.get_session_list()
    return list(ret.values())

if __name__ == "__main__":
    mcp.run(transport='stdio')
    # mcp.run(transport="streamable-http",port = 8000)
    